/* 
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 * 
 * http://www.apache.org/licenses/LICENSE-2.0
 * 
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 *
 */

package org.apache.muse.tools.generator.util;
import java.util.ArrayList;
import java.util.Collection;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Iterator;
import java.util.Map;
import java.util.Set;
import java.util.Map.Entry;

import javax.wsdl.Binding;
import javax.wsdl.BindingFault;
import javax.wsdl.BindingInput;
import javax.wsdl.BindingOperation;
import javax.wsdl.BindingOutput;
import javax.wsdl.Definition;
import javax.wsdl.Fault;
import javax.wsdl.Input;
import javax.wsdl.Message;
import javax.wsdl.Operation;
import javax.wsdl.Output;
import javax.wsdl.Part;
import javax.wsdl.Port;
import javax.wsdl.PortType;
import javax.wsdl.Service;
import javax.wsdl.Types;
import javax.wsdl.WSDLException;
import javax.wsdl.extensions.ExtensionRegistry;
import javax.wsdl.extensions.schema.Schema;
import javax.wsdl.extensions.soap.SOAPAddress;
import javax.wsdl.extensions.soap.SOAPBinding;
import javax.wsdl.extensions.soap.SOAPBody;
import javax.wsdl.extensions.soap.SOAPFault;
import javax.wsdl.extensions.soap.SOAPOperation;
import javax.wsdl.factory.WSDLFactory;
import javax.xml.namespace.QName;

import org.apache.muse.util.messages.Messages;
import org.apache.muse.util.messages.MessagesFactory;
import org.apache.muse.util.xml.XmlUtils;
import org.apache.muse.util.xml.XsdUtils;
import org.apache.muse.ws.addressing.WsaConstants;
import org.apache.muse.ws.resource.properties.WsrpConstants;
import org.apache.muse.ws.wsdl.WsdlUtils;
import org.w3c.dom.Document;
import org.w3c.dom.Element;

import com.ibm.wsdl.extensions.schema.SchemaConstants;
import com.ibm.wsdl.extensions.soap.SOAPConstants;

/**
 * Wraps a WSDL <code>Definition</code> for WsdlMerge. 
 * Also helps in creating the merged WSDL document using
 * wsdl4j.
 * 
 * @author Andrew Eberbach (aeberbac)
 * 
 * @see org.apache.muse.tools.generator.WsdlMerge
 */
public class DefinitionInfo implements DefinitionConstants{

	private static Messages _MESSAGES = MessagesFactory.get(DefinitionInfo.class);
	
	private Definition _definition;

	private PortType _portType;

	private HashMap _operationNameMap = new HashMap();
	
	private WSDLFactory _factory;

	private String _uri;
	
	private String _prefix;
	
	private int _prefixCounter = 0;

	private Element _propertiesElementSequence;

	private Document _document;
	
	private Binding _binding;

	private ExtensionRegistry _registry;

	private Schema _schema;

	private Set _schemaImports = new HashSet();

	private HashMap _schemas = new HashMap();

	public DefinitionInfo(Definition definition, String uri) {
		_uri = uri;		
		_factory = getFactory();
		_registry = _factory.newPopulatedExtensionRegistry();
		
		if (definition == null) {			
			_definition = createDefinition();
			_prefix = addPrefix(_uri);	
			addPrefix(WsrpConstants.NAMESPACE_URI);
			addPrefix(WsaConstants.NAMESPACE_URI);
		} else {
			_definition = definition;			
		}
		
		extractOperations();
	}

	public DefinitionInfo(String uri) {
		this(null, uri);
	}	

	public DefinitionInfo(Definition definition) {
		this(definition, null);
	}

	public Definition getDefinition() {
		return _definition;
	}

	public Collection getOperations() {
		return _operationNameMap.values();
	}
	
	public Types getTypes() {
		return _definition.getTypes();
	}

	private void extractOperations() {

		Map portTypes = _definition.getPortTypes();
		
		if(portTypes.size() != 1) {
			return;
		}
		
		_portType = (PortType) portTypes.values().iterator().next();

		for (Iterator i = _portType.getOperations().iterator(); i.hasNext();) {
			Operation nextOperation = (Operation) i.next();
			String name = nextOperation.getName();
			_operationNameMap.put(name, nextOperation);
		}
	}

	private Operation createOperation(Operation operation) {
		String name = operation.getName();

		Operation newOperation = _definition.createOperation();
		newOperation.setName(name);

		newOperation.setInput(createInput(operation.getInput()));
		if (operation.getOutput() != null){
			newOperation.setOutput(createOutput(operation.getOutput()));
		}

		for (Iterator i = operation.getFaults().values().iterator(); i.hasNext();) {
			newOperation.addFault(createFault((Fault) i.next()));
		}

		newOperation.setUndefined(false);
		
		return newOperation;
	}

	private Fault createFault(Fault fault) {
		Fault newFault = _definition.createFault();
		newFault.setName(fault.getName());

		Message newMessage = createMessage(fault.getMessage());
		_definition.addMessage(newMessage);
		
		newFault.setMessage(newMessage);

		return newFault;
	}

	private Output createOutput(Output output) {
		Output newOutput = _definition.createOutput();
		newOutput.setName(output.getName());

		Message newMessage = createMessage(output.getMessage());
		_definition.addMessage(newMessage);

		newOutput.setMessage(newMessage);
				
		for(Iterator i = output.getExtensionAttributes().entrySet().iterator(); i.hasNext();) {
			Map.Entry next = (Entry) i.next();
			Object value = next.getValue();
			Object key = next.getKey();
			newOutput.setExtensionAttribute((QName)key, value); 
		}
		
		return newOutput;
	}

	private Input createInput(Input input) {
		Input newInput = _definition.createInput();
		newInput.setName(input.getName());

		Message newMessage = createMessage(input.getMessage());
		_definition.addMessage(newMessage);

		newInput.setMessage(newMessage);
		
		for(Iterator i = input.getExtensionAttributes().entrySet().iterator(); i.hasNext();) {
			Map.Entry next = (Entry) i.next();
			Object value = next.getValue();
			Object key = next.getKey();
			newInput.setExtensionAttribute((QName)key, value); 
		}
				
		return newInput;
	}

	private Message createMessage(Message message) {
		Message newMessage = _definition.createMessage();		
		newMessage.setQName(makeLocalQName(message.getQName()));
		
		for (Iterator i = message.getParts().values().iterator(); i.hasNext();) {
			Part nextPart = (Part) i.next();
			
			Part newPart = _definition.createPart();
			newPart.setName(nextPart.getName());
			
			QName qname = nextPart.getElementName();
			if(qname == null) {
				qname = nextPart.getTypeName();
				newPart.setTypeName(qname);
			} else {
				newPart.setElementName(qname);
			}
			
			addPrefix(qname.getNamespaceURI());			
			newMessage.addPart(newPart);
		}

		newMessage.setUndefined(false);
		return newMessage;
	}
	
	private Definition createDefinition() {
		Definition definition = _factory.newDefinition();
		definition.setTargetNamespace(_uri);
		
		PortType portType = definition.createPortType();
		portType.setQName(new QName(_uri, PORT_TYPE));
		portType.setExtensionAttribute(WsrpConstants.DEFAULT_DOCUMENT_QNAME, new QName(_uri,RESOURCE_PROPERTIES));
		portType.setUndefined(false);
		
		definition.addPortType(portType);
				
		_schema = (Schema)createExtension(Types.class, SchemaConstants.Q_ELEM_XSD_2001);
		_schema.setElement(createEmptyProperties());
		
		Types types = definition.createTypes();
		types.addExtensibilityElement(_schema);		
		
		definition.setTypes(types);		
		
		return definition;
	}

	private Element createEmptyProperties() {
		_document = XmlUtils.createDocument();
		Element schema = XmlUtils.createElement(_document, XsdUtils.SCHEMA_QNAME);
		Element element = XmlUtils.createElement(_document, XsdUtils.ELEMENT_QNAME);
		Element complexType = XmlUtils.createElement(_document, XsdUtils.COMPLEX_TYPE_QNAME);
		Element sequence = XmlUtils.createElement(_document, XsdUtils.SEQUENCE_QNAME);
		
		complexType.appendChild(sequence);
		element.appendChild(complexType);
		element.setAttribute(XsdUtils.NAME, RESOURCE_PROPERTIES);
		schema.appendChild(element);
		schema.setAttribute(XmlUtils.TARGET_NS, _uri);
		
		_propertiesElementSequence = sequence;
		
		return schema;
	}

	public void addOperation(Operation operation) {
		String name = operation.getName();
		Operation newOperation = null;
		if (_operationNameMap.containsKey(name)) {
			Object[] filler = { name };
			AbstractCommandLineApp.handleMessage(_MESSAGES.get("DuplicateOperation", filler));
			operation.setName(name);
		} else { 
			newOperation = createOperation(operation);
			_portType.addOperation(newOperation);
			_operationNameMap.put(name, newOperation);
		}
	}

	public void addSchema(Schema schema) {
		String uri = schema.getElement().getAttribute(XmlUtils.TARGET_NS);
		if(uri != null) {

			if(!_schemas.containsKey(uri)) {
				_schemas.put(uri, schema);
				_definition.getTypes().addExtensibilityElement(schema);
			}
		}		
	}

	public void mergeProperties(DefinitionInfo sourceDefinition) {
		Element[] properties = sourceDefinition.getPropertiesDef();
		
		if(properties != null) {
			for(int i=0; i < properties.length; i++) {
				Element property = (Element)_document.importNode(properties[i],true);
				
				String ref = property.getAttribute(XsdUtils.REF);
				QName qname = XmlUtils.parseQName(ref, properties[i]);
								
				String uri = qname.getNamespaceURI();
				String prefix = addPrefix(uri);
				
				property.setAttribute(XsdUtils.REF,makeQName(prefix, qname));
				_propertiesElementSequence.appendChild(property);
				
				addImport(qname);
			}
		}
	}

	private void addImport(QName qname) {
		if (!_schemaImports .contains(qname.getNamespaceURI())) {					
			Element importElement = XmlUtils.createElement(_schema.getElement().getOwnerDocument(), 
					XsdUtils.IMPORT_QNAME);
			importElement.setAttribute(XsdUtils.NAMESPACE, qname.getNamespaceURI());

			Element firstElement = XmlUtils.getFirstElement(_schema.getElement());
			_schema.getElement().insertBefore(importElement, firstElement);
			_schemaImports.add(qname.getNamespaceURI());
		}
	}

	public void createBinding() {
		_binding = _definition.createBinding();
		_binding.setPortType(_portType);
		_binding.setQName(new QName(_uri,BINDING));
		_binding.setUndefined(false);

		SOAPBinding binding = (SOAPBinding)createExtension(Binding.class, SOAPConstants.Q_ELEM_SOAP_BINDING);
		binding.setStyle(DOCUMENT_STYLE);

		addPrefix(HTTP_SOAP_URI);
		addPrefix(SOAP_WSDL_URI);
		binding.setTransportURI(HTTP_SOAP_URI);
		
		_binding.addExtensibilityElement(binding);			
		_definition.addBinding(_binding);
		
		for(Iterator i = _portType.getOperations().iterator(); i.hasNext(); ) {
			BindingOperation bindingOperation = createOperationBinding((Operation)i.next());
			_binding.addBindingOperation(bindingOperation);
		}			
	}
	
	public BindingOperation createOperationBinding(Operation operation) {
		BindingOperation bindingOperation = _definition.createBindingOperation();
		bindingOperation.setName(operation.getName());

		SOAPOperation soapOperation = (SOAPOperation)createExtension(BindingOperation.class, SOAPConstants.Q_ELEM_SOAP_OPERATION);
		bindingOperation.addExtensibilityElement(soapOperation);
		
		bindingOperation.setBindingInput(createInputBinding(operation.getInput()));
		if (operation.getOutput() != null){
			bindingOperation.setBindingOutput(createOutputBinding(operation.getOutput()));
		}

		for(Iterator i = operation.getFaults().values().iterator(); i.hasNext();) {
			bindingOperation.addBindingFault(createFaultBinding((Fault)i.next()));
		}
				
		return bindingOperation;
	}

	private BindingFault createFaultBinding(Fault fault) {
		BindingFault bindingFault = _definition.createBindingFault();
		bindingFault.setName(fault.getName());
		
		SOAPFault soapFault = (SOAPFault)createExtension(BindingFault.class, SOAPConstants.Q_ELEM_SOAP_FAULT);
		soapFault.setUse(LITERAL_USE);
		soapFault.setName(fault.getName());
		bindingFault.addExtensibilityElement(soapFault);

		return bindingFault;
	}

	private BindingOutput createOutputBinding(Output output) {
		BindingOutput bindingOutput = _definition.createBindingOutput();
		bindingOutput.setName(output.getName());
		
		SOAPBody soapBody = (SOAPBody)createExtension(BindingOutput.class, SOAPConstants.Q_ELEM_SOAP_BODY);
		soapBody.setUse(LITERAL_USE);
		
		ArrayList list = new ArrayList();
		list.add(ENCODING_URI);		
		soapBody.setEncodingStyles(list);
		
		bindingOutput.addExtensibilityElement(soapBody);
		return bindingOutput;
	}

	private BindingInput createInputBinding(Input input) {
		BindingInput bindingInput = _definition.createBindingInput();
		bindingInput.setName(input.getName());
		
		SOAPBody soapBody = (SOAPBody) createExtension(BindingInput.class, SOAPConstants.Q_ELEM_SOAP_BODY);
		soapBody.setUse(LITERAL_USE);
		ArrayList list = new ArrayList();
		list.add(ENCODING_URI);
		soapBody.setEncodingStyles(list);
		bindingInput.addExtensibilityElement(soapBody);

		return bindingInput;
	}

	public void createService(String location) {
		Service service = _definition.createService();
		service.setQName(new QName(_uri,SERVICE));
		Port port = _definition.createPort();
		port.setBinding(_binding);
		port.setName(PORT);
		service.addPort(port);

		SOAPAddress soapAddress = (SOAPAddress)createExtension(Port.class, SOAPConstants.Q_ELEM_SOAP_ADDRESS);
		soapAddress.setLocationURI(location);
		port.addExtensibilityElement(soapAddress);
				
		_definition.addService(service);
	}

	private String addPrefix(String uri) {
		String prefix = getPrefix(uri);		
		_definition.addNamespace(prefix, uri);		
		return prefix;
	}
	
	private Object createExtension(Class targetClass, QName qname) {
		try {
			return _registry.createExtension(targetClass, qname);
		} catch (WSDLException e) {
			AbstractCommandLineApp.handleErrorAndExit(_MESSAGES.get("WsdlExtensionFailed"), e);
		}
		return null;
	}
	
	private QName makeLocalQName(QName name) {
		return new QName(_uri, name.getLocalPart(), _prefix);
	}

	private String getPrefix(String uri) {
		String prefix = _definition.getPrefix(uri);	
		if(prefix == null) {
			prefix =  PREFIX + _prefixCounter++;
		} 
		return prefix;
	}
	
	private static WSDLFactory getFactory() {
		try {
			return WSDLFactory.newInstance();
		} catch (WSDLException e) {
			throw new RuntimeException("Could not get factory");
		}				
	}

	private String makeQName(String prefix, QName qname) {
		return prefix  + ":" + qname.getLocalPart();
	}
	
	private Element[] getPropertiesDef() {
		if(_portType == null) {
			return null;
		}
		Object propertiesDef = _portType.getExtensionAttribute(WsrpConstants.DEFAULT_DOCUMENT_QNAME);
		if(propertiesDef != null && propertiesDef instanceof QName) {
			QName qname = (QName)propertiesDef;
			Element propertiesElement = find(qname);	
			
			if(propertiesElement != null) {
				return XmlUtils.findInSubTree(propertiesElement, XsdUtils.ELEMENT_QNAME);
			}
		}
		return null;
	}

	private Element find(QName qname) {
		Types types = _definition.getTypes();
		
		for(Iterator i = types.getExtensibilityElements().iterator(); i.hasNext();) {
			Object next = i.next();
			if(next instanceof Schema) {
				Schema nextSchema = (Schema)next;
				Element type = WsdlUtils.getElementDeclaration(nextSchema.getElement(), qname);
				if(type != null) {
					return type;
				}
			}
		}
		return null;
	}
}